import paddle as pp
from paddle import nn
from deberta.deberta_model import DebertaModel
from deberta.context_pooler import ContextPooler
from deberta.dropout import StableDropout


class DebertaForSequenceClassification(nn.Layer):
    def __init__(self, config):
        super().__init__()
        self.config = config

        num_labels = getattr(config, "num_labels", 3)  # for mnli
        self.num_labels = num_labels

        self.deberta = DebertaModel(config)
        self.pooler = ContextPooler(config)
        output_dim = self.pooler.output_dim

        self.classifier = nn.Linear(output_dim, num_labels)
        drop_out = getattr(config, "cls_dropout", None)
        drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
        self.dropout = StableDropout(drop_out)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        outputs = self.deberta(
            input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )

        encoder_layer = outputs['last_hidden_state'] if isinstance(outputs, dict) and 'last_hidden_state' in outputs else outputs[0]
        pooled_output = self.pooler(encoder_layer)
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        return dict(
            logits=logits,
            hidden_states=outputs['hidden_states'] if 'hidden_states' in outputs else None,
            attentions=outputs['attentions'] if 'attentions' in outputs else None
        ) if return_dict else logits

